package jmasswiz;

import com.compomics.util.Util;
import com.compomics.util.experiment.biology.EnzymeFactory;
import com.compomics.util.experiment.identification.SequenceFactory;
import com.compomics.util.experiment.massspectrometry.SpectrumFactory;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.ArrayList;
import javax.swing.JOptionPane;
import uk.ac.ebi.jmzml.xml.io.MzMLUnmarshallerException;
import com.compomics.util.experiment.massspectrometry.MSnSpectrum;
import com.compomics.util.protein.Header;
import com.compomics.util.experiment.biology.Protein;
import com.compomics.util.experiment.biology.Enzyme;
import com.compomics.util.experiment.biology.Ion;
import com.compomics.util.experiment.biology.IonFactory;
import com.compomics.util.experiment.biology.PTM;
import com.compomics.util.experiment.biology.PTMFactory;
import com.compomics.util.experiment.biology.Peptide;
import com.compomics.util.experiment.biology.ions.PeptideFragmentIon;
import com.compomics.util.experiment.identification.NeutralLossesMap;
import com.compomics.util.experiment.identification.SpectrumAnnotator;
import com.compomics.util.experiment.identification.matches.IonMatch;
import com.compomics.util.experiment.identification.matches.ModificationMatch;
import java.util.HashMap;

/**
 * A demo class showing how to use the compomics utilities library.
 *
 * @author Harald Barsnes
 * @author Marc Vaudel
 */
public class jMassWiz {

    /**
     * The spectrum factory for reading the spectra from the mgf file.
     */
    private SpectrumFactory spectrumFactory = SpectrumFactory.getInstance();
    /**
     * The sequence factory for reading the sequence from the FASTA file.
     */
    private SequenceFactory sequenceFactory = SequenceFactory.getInstance(100000);
    /**
     * The compomics enzyme factory.
     */
    private EnzymeFactory enzymeFactory = EnzymeFactory.getInstance();
    /**
     * The XML file containing the enzymes.
     */
    public static final String ENZYME_FILE = "resources/conf/enzymes.xml";
    /**
     * The XML file containing the default modifications.
     */
    public static final String DEFAULT_MODIFICATIONS_FILE = "resources/conf/mods.xml";
    /**
     * The compomics PTM factory.
     */
    private PTMFactory ptmFactory = PTMFactory.getInstance();

    /**
     * A main class showing simple uses of compomics utilities.
     *
     * @param args the command line arguments
     */
    public static void main(String[] args) {

        
        // check if two input files are provided
        if (args.length < 2 || args.length > 2) {
            JOptionPane.showMessageDialog(null,
                    "Incorrect number of command line arguments!\n"
                    + "Should be two: 1) mgf file + 2) FASTA file", "Input Error", JOptionPane.WARNING_MESSAGE);
            System.exit(0);
        } else {
            // set up a jMassWiz object
            jMassWiz jMassWiz = new jMassWiz();

            // parse the mgf file
            jMassWiz.parseSpectra(new File(args[0]));

            // Annotate the first spectrum of the mgf file with a dummy peptide
            String fileName = Util.getFileName(args[0]);
            String spectrumTitle = SpectrumFactory.getInstance().getSpectrumTitles(fileName).get(0);
            jMassWiz.annotateSpectrum(fileName, spectrumTitle);

            // parse the fasta file
            jMassWiz.parseFastaFile(new File(args[1]));
        }
    }

    /**
     * Parse the spectra and print some example details to screen.
     *
     * @param mgfFile an mgf file
     */
    private void parseSpectra(File mgfFile) {

        // check if the file exists
        if (mgfFile.exists()) {
            try {
                // add all the spectra to the factory
                spectrumFactory.addSpectra(mgfFile);

                // get the number of spectra
                int numberOfSpectra = spectrumFactory.getNSpectra();
                System.out.println("\nNumber of spectra in mgf file: " + numberOfSpectra + "\n");

                // get the titles of all the spectra in the mgf file
                ArrayList<String> allSpectrumTitles = spectrumFactory.getSpectrumTitles(mgfFile.getName());

                // iterate the spectra and print the spectrum title and precursor mz, break after 5 spectra
                int spectrumCount = 0;
                for (String spectrumTitle : allSpectrumTitles) {
                    MSnSpectrum tempSpectrum = (MSnSpectrum) spectrumFactory.getSpectrum(mgfFile.getName(), spectrumTitle);
                    double precursorMz = tempSpectrum.getPrecursor().getMz();
                    System.out.println("spectrumTitle: " + spectrumTitle + "\t" + precursorMz);
                    spectrumCount++;
                    if (spectrumCount > 5) {
                        break;
                    }
                }

                // print an empty line
                System.out.println();

            } catch (FileNotFoundException ex) {
                ex.printStackTrace();
            } catch (IOException ex) {
                ex.printStackTrace();
            } catch (ClassNotFoundException ex) {
                ex.printStackTrace();
            } catch (MzMLUnmarshallerException ex) {
                ex.printStackTrace();
            }
        }
    }

    /**
     * Parse the FASTA file.
     *
     * @param fastaFile a FASTA file
     */
    private void parseFastaFile(File fastaFile) {

        // check if the file exists
        if (fastaFile.exists()) {
            try {
                // add all the sequences to the factory
                sequenceFactory.loadFastaFile(fastaFile);

                // get the number of sequences
                int numberOfSequences = sequenceFactory.getNTargetSequences();
                System.out.println("\nNumber of proteins FASTA file: " + numberOfSequences + "\n");

                // get all the protein accession numbers
                ArrayList<String> allAccessions = sequenceFactory.getAccessions();

                // print the first five accession numbers
                int proteinCount = 0;
                for (String accession : allAccessions) {
                    System.out.println("accession: " + accession);
                    proteinCount++;
                    if (proteinCount > 5) {
                        break;
                    }
                }

                // print an empty line
                System.out.println();

                // get the fasta header of the first protein
                String firstAccessionNumber = allAccessions.get(0);
                Header fastaHeader = sequenceFactory.getHeader(firstAccessionNumber);

                // print the protein accession, description and db type
                System.out.println("accession: " + fastaHeader.getAccession() + ", description: " + fastaHeader.getDescription() + ", db type: " + fastaHeader.getDatabaseType() + "\n");

                // get the protein sequence
                Protein protein = sequenceFactory.getProtein(firstAccessionNumber);
                String proteinSequence = protein.getSequence();
                System.out.println("proteinSequence:\n" + proteinSequence + "\n");

                // load the enzymes
                loadEnzymes();

                // digest the protein using trypsin, 2 missed cleavages, and peptide lenghts of 5-15
                int nMissedCleavages = 1;
                int minPeptideLength = 5;
                int maxPeptideLength = 15;
                Enzyme trypsin = enzymeFactory.getEnzyme("Trypsin");
                ArrayList<String> allPeptides = trypsin.digest(proteinSequence, nMissedCleavages, minPeptideLength, maxPeptideLength);

                // iterate and print the peptides
                System.out.println("peptides: ");
                for (String peptide : allPeptides) {
                    System.out.println(peptide);
                }

                // print an empty line
                System.out.println();

            } catch (FileNotFoundException ex) {
                ex.printStackTrace();
            } catch (IOException ex) {
                ex.printStackTrace();
            } catch (ClassNotFoundException ex) {
                ex.printStackTrace();
            } catch (IllegalArgumentException ex) {
                ex.printStackTrace();
            } catch (InterruptedException ex) {
                ex.printStackTrace();
            }
        }
    }

    /**
     * Annotate the given spectrum.
     *
     * @param spectrumFile the spectrum file
     * @param spectrumTitle the title of the spectrum to annotate
     */
    private void annotateSpectrum(String spectrumFile, String spectrumTitle) {

        try {
            // create a new peptide, sequence PEPTIDEK with a phosphorylation on T(4), belonging to the proteins A and B. 
            // note that the modification is indexed by its name in the PTMFactory.
            // proteins are indexed by their accession, given by the fasta file
            String sequence = "PEPTIDEK";
            ArrayList<ModificationMatch> modificationMatches = new ArrayList<ModificationMatch>();
            modificationMatches.add(new ModificationMatch("phosphorylation of t", true, 4));
            ArrayList<String> proteins = new ArrayList<String>();
            proteins.add("A");
            proteins.add("B");
            Peptide peptide = new Peptide(sequence, proteins, modificationMatches);
            System.out.println("Peptide " + peptide.getSequence() + " created.");

            // the PTM details can be obtained from the PTM factory
            loadModifications();
            PTM ptm = ptmFactory.getPTM("phosphorylation of t");
            System.out.println("\nModification " + ptm.getName() + " is of mass " + ptm.getMass() + ".\n");

            // the IonFactory allows the estimation of the theoretic fragment ions
            // of course you can fine tune the ion selection with regard to charge and neutral losses
            ArrayList<Ion> fragments = IonFactory.getInstance().getFragmentIons(peptide);

            // print the first ten fragment ions
            int fragmentIonCounter = 0;
            for (Ion ion : fragments) {
                System.out.println("Theoretic fragment ion " + ion.getName() + " of mass " + ion.getTheoreticMass() + ".");
                fragmentIonCounter++;
                if (fragmentIonCounter > 10) {
                    break;
                }
            }

            // print an empty line
            System.out.println();

            // using the SpectrumAnnotator you can match these ions in a spectrum and collect ion matches
            SpectrumAnnotator spectrumAnnotator = new SpectrumAnnotator();

            // we will annotate the b and y ions only, all implemented ions can be found in com.compomics.util.experiment.biology.ions
            HashMap<Ion.IonType, ArrayList<Integer>> iontypes = new HashMap<Ion.IonType, ArrayList<Integer>>();
            ArrayList<Integer> fragmentIonType = new ArrayList<Integer>();
            fragmentIonType.add(PeptideFragmentIon.B_ION);
            fragmentIonType.add(PeptideFragmentIon.Y_ION);
            iontypes.put(Ion.IonType.PEPTIDE_FRAGMENT_ION, fragmentIonType);

            // the neutral losses can be automatically deduced from the peptide sequence
            NeutralLossesMap neutralLossesMap = SpectrumAnnotator.getDefaultLosses(peptide);

            // we look for charge 1 fragments only
            ArrayList<Integer> charges = new ArrayList<Integer>();
            charges.add(1);

            // let's say a precursor charge of 2
            int precursorCharge = 2;

            // the spectrum to annotate
            MSnSpectrum spectrum = (MSnSpectrum) spectrumFactory.getSpectrum(spectrumFile, spectrumTitle);

            // the intensity limit to use, here 0 and the m/z tolerance, here 0.5
            double intensityLimit = 0;
            double mzTolerance = 0.5;

            // rhe result is a list of ion matches
            ArrayList<IonMatch> matches = spectrumAnnotator.getSpectrumAnnotation(
                    iontypes, neutralLossesMap, charges, precursorCharge,
                    spectrum, peptide, intensityLimit, mzTolerance, false);

            if (matches.isEmpty()) {
                System.out.println("No ions matched!");
            }
            for (IonMatch ionMatch : matches) {
                System.out.println("Ion matched at m/z: " + ionMatch.peak.mz);
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * Loads the enzymes from the enzyme file into the enzyme factory.
     */
    private void loadEnzymes() {
        try {
            enzymeFactory.importEnzymes(new File(ENZYME_FILE));
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * Load modifications from the default list. Note that user modifications
     * can be added/edited using com.compomics.util.gui.ptm.ModificationsDialog
     * Modifications are stored in the user folder and compatible across
     * versions.
     */
    private void loadModifications() {
        try {
            File modsFile = new File(DEFAULT_MODIFICATIONS_FILE);
            ptmFactory.importModifications(modsFile, false);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
